import os
import argparse
import time
from core.deploy import MiniCPMVDeploy
from inference.image_inference import ImageInferencer
from inference.video_inference import VideoInferencer
from utils.batch_processor import process_directory

class MultiInferServer:
    """
    MiniCPM-V 多任务推理服务器，实现模型一次加载，多次使用
    """
    def __init__(self, config_file=None):
        # 初始化部署类并加载模型（只加载一次）
        self.deploy = MiniCPMVDeploy(config_file)
        self.model_loaded = False
        self.start_time = time.time()
        
    def load_model(self):
        """
        加载模型（只需要执行一次）
        """
        if not self.model_loaded:
            print("\n===== 开始加载模型 =====")
            start_time = time.time()
            if self.deploy.load_model():
                self.model_loaded = True
                load_time = time.time() - start_time
                print(f"模型加载成功，耗时: {load_time:.2f}秒")
                print("模型已加载至内存，可进行多次推理任务")
            else:
                print("模型加载失败，请检查模型配置和路径")
        else:
            print("模型已经加载，无需重复加载")
        
    def run_batch_image_infer(self, image_dir=None, prompt=None, output_dir=None, 
                             export_csv=None, max_workers=None, recursive=False, 
                             no_save=False):
        """
        运行批量图片推理
        """
        if not self.model_loaded:
            print("错误: 模型尚未加载，请先加载模型")
            return False
        
        # 使用配置文件中的默认值
        if not image_dir:
            image_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'images')
        
        # 验证图片目录是否存在
        if not os.path.exists(image_dir):
            print(f"错误: 目录不存在: {image_dir}")
            return False
        
        # 如果没有提供提示词，使用默认提示词
        if not prompt:
            prompt = "请详细描述这张图片的内容"
        
        # 更新输出目录配置（如果提供）
        if output_dir:
            self.deploy.update_config({"paths": {"output_dir": output_dir}})
        
        print("\n===== 开始批量图片推理 =====")
        print(f"图片目录: {image_dir}")
        print(f"提示词模板: {prompt}")
        print(f"是否递归搜索: {recursive}")
        if max_workers is not None:
            print(f"最大工作线程数: {max_workers}")
        if export_csv:
            print(f"CSV导出路径: {export_csv}")
        print(f"是否保存单个结果文件: {not no_save}")
        
        # 开始批量处理
        start_time = time.time()
        process_directory(
            dir_path=image_dir,
            inferencer_class=ImageInferencer,
            prompt=prompt,
            save_result=not no_save,
            export_csv=export_csv,
            max_workers=max_workers,
            recursive=recursive,
            deploy=self.deploy  # 传递已加载的deploy实例
        )
        
        infer_time = time.time() - start_time
        print(f"\n批量图片推理完成，耗时: {infer_time:.2f}秒")
        print("模型仍在内存中，可继续执行其他推理任务")
        return True
        
    def run_batch_video_infer(self, video_dir=None, prompt=None, output_dir=None, 
                             export_csv=None, max_workers=None, recursive=False, 
                             no_save=False):
        """
        运行批量视频推理
        """
        if not self.model_loaded:
            print("错误: 模型尚未加载，请先加载模型")
            return False
        
        # 使用配置文件中的默认值
        if not video_dir:
            video_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'videos')
        
        # 验证视频目录是否存在
        if not os.path.exists(video_dir):
            print(f"错误: 目录不存在: {video_dir}")
            return False
        
        # 如果没有提供提示词，使用默认提示词
        if not prompt:
            prompt = "请详细描述这个视频的内容"
        
        # 更新输出目录配置（如果提供）
        if output_dir:
            self.deploy.update_config({"paths": {"output_dir": output_dir}})
        
        print("\n===== 开始批量视频推理 =====")
        print(f"视频目录: {video_dir}")
        print(f"提示词模板: {prompt}")
        print(f"是否递归搜索: {recursive}")
        if max_workers is not None:
            print(f"最大工作线程数: {max_workers}")
        if export_csv:
            print(f"CSV导出路径: {export_csv}")
        print(f"是否保存单个结果文件: {not no_save}")
        
        # 开始批量处理
        start_time = time.time()
        process_directory(
            dir_path=video_dir,
            inferencer_class=VideoInferencer,
            prompt=prompt,
            save_result=not no_save,
            export_csv=export_csv,
            max_workers=max_workers,
            recursive=recursive,
            deploy=self.deploy  # 传递已加载的deploy实例
        )
        
        infer_time = time.time() - start_time
        print(f"\n批量视频推理完成，耗时: {infer_time:.2f}秒")
        print("模型仍在内存中，可继续执行其他推理任务")
        return True
        
    def unload_model(self):
        """
        卸载模型，释放内存
        """
        if self.model_loaded:
            self.deploy.unload_model()
            self.model_loaded = False
            print("模型已卸载，内存已释放")
        else:
            print("模型尚未加载，无需卸载")
            
    def show_stats(self):
        """
        显示服务器统计信息
        """
        uptime = time.time() - self.start_time
        print("\n===== 服务器统计信息 =====")
        print(f"运行时间: {uptime:.2f}秒")
        print(f"模型加载状态: {'已加载' if self.model_loaded else '未加载'}")
        if self.model_loaded:
            device_info = self.deploy.get_device_info()
            print(f"设备信息: {device_info}")
        print("======================")

    def start_interactive_mode(self):
        """
        启动交互式模式
        """
        print("\n===== MiniCPM-V 多任务推理服务器 =====")
        print("模型一次加载，多次使用 - 交互式模式")
        print("输入 'help' 获取可用命令")
        print("==================================")
        
        while True:
            cmd = input("\n请输入命令: ").strip().lower()
            
            if cmd in ['exit', 'quit', 'q']:
                self.unload_model()
                print("感谢使用MiniCPM-V多任务推理服务器！")
                break
            
            elif cmd in ['help', 'h']:
                self.show_help()
            
            elif cmd == 'load':
                self.load_model()
            
            elif cmd == 'unload':
                self.unload_model()
            
            elif cmd == 'stats':
                self.show_stats()
            
            elif cmd.startswith('image'):
                # 解析参数 - 简单实现，可以扩展为完整的参数解析
                parts = cmd.split(' ')
                image_dir = parts[1] if len(parts) > 1 else None
                self.run_batch_image_infer(image_dir=image_dir)
            
            elif cmd.startswith('video'):
                # 解析参数 - 简单实现，可以扩展为完整的参数解析
                parts = cmd.split(' ')
                video_dir = parts[1] if len(parts) > 1 else None
                self.run_batch_video_infer(video_dir=video_dir)
            
            else:
                print(f"未知命令: {cmd}")
                print("输入 'help' 获取可用命令")
                
    def show_help(self):
        """
        显示帮助信息
        """
        print("\n===== 可用命令 =====")
        print("load              - 加载模型（只需执行一次）")
        print("unload            - 卸载模型，释放内存")
        print("stats             - 显示服务器统计信息")
        print("image [目录]      - 运行批量图片推理，可选指定图片目录")
        print("video [目录]      - 运行批量视频推理，可选指定视频目录")
        print("help, h           - 显示帮助信息")
        print("exit, quit, q     - 退出服务器并卸载模型")
        print("=================")


def main():
    """
    多任务推理服务器主函数
    """
    print("===== MiniCPM-V 多任务推理服务器 ======")
    print("基于模块化架构，实现模型一次加载，多次使用")
    print("====================================")
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='MiniCPM-V 多任务推理服务器')
    parser.add_argument('--config', type=str, default=None,
                        help='配置文件路径')
    parser.add_argument('--interactive', action='store_true',
                        help='以交互式模式运行')
    parser.add_argument('--image_dir', type=str, default=None,
                        help='图片文件夹路径，用于非交互式模式')
    parser.add_argument('--video_dir', type=str, default=None,
                        help='视频文件夹路径，用于非交互式模式')
    args = parser.parse_args()
    
    # 创建服务器实例
    server = MultiInferServer(args.config)
    
    # 根据模式选择运行方式
    if args.interactive or (not args.image_dir and not args.video_dir):
        # 交互式模式
        server.start_interactive_mode()
    else:
        # 非交互式模式
        # 加载模型
        server.load_model()
        
        # 运行指定的推理任务
        if args.image_dir:
            server.run_batch_image_infer(image_dir=args.image_dir)
        elif args.video_dir:
            server.run_batch_video_infer(video_dir=args.video_dir)
        
        # 任务完成后卸载模型
        server.unload_model()
        
        print("\n===== 多任务推理服务器完成 ====")
        print("感谢使用MiniCPM-V多任务推理服务器！")


if __name__ == "__main__":
    main()